#include "coroutine.h"
#include "../coro_call.h"

namespace coroutine {

/// The single instance for 'Scheduler' per thread
thread_local CoroSched_* instance_ = nullptr;
/// Current running handle
thread_local CoroHandle curHandle_ = INVALID_CORO_HANDLE;

/// Helper function to switch fiber in windows platform
static inline void switchToFiber(CoroHandle handle) {
    if (curHandle_ == handle) {
        return;
    }
    curHandle_ = handle;
    ::SwitchToFiber(handle);
}

CoroSched_& CoroSched_::instance() {
    if (!instance_) {        
        instance_ = new CoroSched_();
    }
    return *instance_;
}

static void destroyCoros_(CoroList& list, CoroMap& coros) {
    for (auto id : list) {
        auto it = coros.find(id);
        if (it != coros.end()) {
            auto* coro = it->second;
            coros.erase(it);
            delete coro;
        }
    }
    list.clear();
}

void CoroSched_::shutdown() {
    if (instance_) {
        instance_->closeAll();
        delete instance_;
        instance_ = nullptr;
    }
}

CoroID CoroSched_::current() {
    if (current_) {
        if (curHandle_ != mainCoro_) {
            return current_->getID();
        } else {
            return 1;
        }
    } else {
        return 1;
    }
}

void CoroSched_::closeAll() {
    destroyCoros_(coroDead_, coros_);
    destroyCoros_(coroSuspend_, coros_);
    destroyCoros_(coroReady_, coros_);
    remoteCall_->clear();
    timer_.clear();
}

void CoroSched_::push(CoroID id, std::any&& msg, bool wakeup) {
    auto it = coros_.find(id);
    if (it == coros_.end()) {
        return;
    }
    auto* coro = it->second;
    if (coro->getStatus() == Status_::DEAD) {
        return;
    }
    coro->queue().emplace_back(Mail(id, std::forward<std::any>(msg)));
    // TODO pending detail reason
    if (coro->getStatus() == Status_::SUSPEND) {
        if (coro->checkAndReset(WaitStatus_::MSG_WAIT)) {
            resume(id);
            return;
        }
    }
    if (wakeup) {
        resume(id);
    }
}

Mail CoroSched_::pop(CoroID id) {
    auto it = coros_.find(id);
    if (it == coros_.end()) {
        return Mail();
    }
    auto& queue = it->second->queue();
    if (!queue.empty()) {
        Mail mail(std::move(queue.front()));
        queue.pop_front();
        return std::move(mail);
    }
    return Mail();
}

Mail CoroSched_::pop_wait(CoroID id) {
    auto it = coros_.find(id);
    if (it == coros_.end()) {
        throw std::runtime_error("Coroutine not found");
    }
    auto* coro = it->second;
    auto& queue = it->second->queue();
    if (queue.empty()) {
        coro->waitFor(WaitStatus_::MSG_WAIT);
        yield();
    }
    if (!queue.empty()) {
        Mail mail(std::move(queue.front()));
        queue.pop_front();
        return std::move(mail);
    }
    return Mail();
}

Mail CoroSched_::pop_wait(CoroID id, std::time_t milliSeconds) {
    auto it = coros_.find(id);
    if (it == coros_.end()) {
        throw std::runtime_error("Coroutine not found");
    }
    auto* coro = it->second;
    auto& queue = it->second->queue();
    if (queue.empty()) {
        auto timerHandle = timer_.addTimerOnce(nullptr, milliSeconds,
            static_cast<decltype(TimerNode_::user_data)>(id));
        coro->setTimerHandle(timerHandle);
        coro->waitFor(WaitStatus_::MSG_WAIT);
        yield();
    }
    if (!queue.empty()) {
        Mail mail(std::move(queue.front()));
        queue.pop_front();
        return std::move(mail);
    }
    return Mail();
}

void CoroSched_::sleep(CoroID id, std::time_t milliSeconds) {
    auto it = coros_.find(id);
    if (it == coros_.end()) {
        throw std::runtime_error("Coroutine not found");
    }
    auto* coro = it->second;
    auto timerHandle = timer_.addTimerOnce(nullptr, milliSeconds,
        static_cast<decltype(TimerNode_::user_data)>(id));
    coro->setTimerHandle(timerHandle);
    coro->waitFor(WaitStatus_::SLEEP);
    yield();
}

CoroSched_::CoroSched_() {
    // First of all we must switch current windows main thread to windows fiber
    // and then, we'll switch main fiber to another child fiber
	if (curHandle_ == INVALID_CORO_HANDLE) {
		mainCoro_ = ::ConvertThreadToFiber(nullptr);
	} else {
		mainCoro_ = curHandle_;
	}
    current_ = nullptr;
    remoteCall_ = new Call();
}

CoroSched_::~CoroSched_() {
    if (mainCoro_) {
        curHandle_ = mainCoro_;
    } else {
		curHandle_ = INVALID_CORO_HANDLE;
	}
    if (remoteCall_) {
        delete remoteCall_;
    }
    mainCoro_ = INVALID_CORO_HANDLE;
    current_ = nullptr;
    remoteCall_ = nullptr;
}

bool CoroSched_::sched() {
    while (true) {
        auto handle = timer_.runOnce();
        if (!handle) {
            break;
        }
        auto it = coros_.find(static_cast<CoroID>(handle->user_data));
        if (it == coros_.end()) {
            continue;
        }
        auto* coro = it->second;
        if (coro->checkAndReset(WaitStatus_::SLEEP)) {
            coro->setTimerHandle(nullptr);
            // wake up sleepy coroutine
            switchTo_(coro, Status_::READY);
        } else if (coro->checkAndReset(WaitStatus_::MSG_WAIT)) {
            coro->setTimerHandle(nullptr);
            // wake up sleepy coroutine
            switchTo_(coro, Status_::READY);
        }
    }
    do {
        current_ = nullptr;
        // Cleanup all dead coroutine
        destroyCoros_(instance_->coroDead_, instance_->coros_);
        if (!current_ && !coroReady_.empty()) {
            // We have a coroutine in Status::READY to run
            current_ = coros_[coroReady_.front()->data];
        }
        if (current_) {
            // Now we have a coroutine to run
            switchToFiber(current_->getHandle());
        }
    } while (current_);
    return true;
}

void __stdcall CoroSched_::entry_(LPVOID lpParam) {
    auto* coro = reinterpret_cast<Coro_*>(lpParam);
    if (coro->getEntry()) {
        // Call user defined entry method
        coro->getEntry()(coro->getUsrPtr());
    }
    // Switch from Status::Ready to Status::DEAD
    CoroSchedRef.switchTo_(coro, Status_::DEAD);
}

void CoroSched_::switchTo_(Coro_* coro, Status_ status) {    
    switch (status) {
        case Status_::DEAD:
            if (coro->getStatus() == Status_::READY) {
                // READY to DEAD
                coroDead_.push_back(coroReady_.remove(coro->node()));
                coro->setStatus(status);
            } else if (coro->getStatus() == Status_::SUSPEND) {
                // SUSPEND to DEAD
                coroDead_.push_back(coroSuspend_.remove(coro->node()));
                coro->setStatus(status);
            } else if (coro->getStatus() == Status_::NONE) {
                // NONE to DEAD
                coroDead_.push_back(coro->node());
                coro->setStatus(status);
            } else {                
                // DEAD to DEAD
            }
            current_ = nullptr;
            break;
        case Status_::SUSPEND:
            if (coro->getStatus() == Status_::READY) {
                // READY to SUSPEND
                coroSuspend_.push_back(coroReady_.remove(coro->node()));
                coro->setStatus(status);
            } else {
                // NONE to SUSPEND
                // SUSPEND to SUSPEND
                // DEAD to SUSPEND
            }
            current_ = nullptr;
            break;
        case Status_::READY:
            if (coro->getStatus() == Status_::NONE) {
                // NONE to READY, use spawn method
                auto* node = coroReady_.add_tail_node(coro->getID());
                coro->setNode(node);
                coro->setStatus(status);
                coros_.emplace(coro->getID(), coro);
                return;
            } else if (coro->getStatus() == Status_::SUSPEND) {
                // SUSPEND to READY
                coroReady_.push_back(coroSuspend_.remove(coro->node()));
                coro->setStatus(status);
                // Clean all wait condition
                coro->resetWait();
                if (!coro->isCancel()) {
                    current_ = coros_[coroReady_.back()->data];
                } else {
                    current_ = coros_[coroReady_.front()->data];
                }
            } else {
                // READY to READY
                // DEAD to READY
                current_ = nullptr;
            }
            break;
        default:
            break;
    }
    if (current_) {
        switchToFiber(current_->getHandle());
    } else {
        // Yield to main coroutine
        switchToFiber(mainCoro_);
    }
}

void CoroSched_::resume(CoroID id) {
    if (id == INVALID_CORO_ID) {
        return;
    }
    if (id == 1) {
        switchToFiber(mainCoro_);
        return;
    }
    auto it = coros_.find(id);
    if (it != coros_.end()) {
        if (!coro_is_main()) {
            // switch current to SUSPEND
            coroSuspend_.push_back(coroReady_.remove(current_->node()));
            current_->resetWait();
            current_->waitFor(WaitStatus_::YIELD);
            current_->setStatus(Status_::SUSPEND);
        }
        // Wake up a suspend coroutine and yield cpu to it
        auto* coro = it->second;
        if (coro->getStatus() == Status_::SUSPEND) {
            if (!coro->isCancel() && coro->isWaitBy(WaitStatus_::SLEEP)) {
                // Sleepy ONLY can be waked up by timer
                return;
            }
            // For other wait status, except WaitStatus_::SLEEP
            switchTo_(coro, Status_::READY);
        } else if (coro->getStatus() == Status_::READY) {
            // READY
            current_ = coro;
            switchToFiber(coro->getHandle());
        } else {
            // NONE, DEAD
        }
    }
}

void CoroSched_::wakeup(CoroID id) {
    if (id == INVALID_CORO_ID) {
        return;
    }
    if (id == 1) {
        switchToFiber(mainCoro_);
        return;
    }
    auto it = coros_.find(id);
    if (it != coros_.end()) {
        // Wake up a suspend coroutine and yield cpu to it
        auto* coro = it->second;
        if (coro->getStatus() == Status_::SUSPEND) {
            if (!coro->isCancel() && coro->isWaitBy(WaitStatus_::SLEEP)) {
                // Sleepy ONLY can be waked up by timer
                return;
            }
            // For other wait status, except WaitStatus_::SLEEP
            switchTo_(coro, Status_::READY);
        } else if (coro->getStatus() == Status_::READY) {
            // READY
            current_ = coro;
            switchToFiber(coro->getHandle());
        } else {
            // NONE, DEAD
        }
    }
}

void CoroSched_::yield() {
    // Suspends myself and switch to main coroutine
    if (current_) {
        if (current_->isCancel()) {
            throw CoroCancelException("Coroutine Cancellation");
        }
        if (!current_->waitFlag_.any()) {
            current_->waitFor(WaitStatus_::YIELD);
        }
        switchTo_(current_, Status_::SUSPEND);
        if (current_->isCancel()) {
            throw CoroCancelException("Coroutine Cancellation");
        }
    } else {
        sched();
    }
}

void CoroSched_::yield(CoroID id) {
    auto it = coros_.find(id);
    if (it == coros_.end()) {
        return;
    }
    auto* coro = it->second;
    if (!coro->isWaitBy(WaitStatus_::YIELD)) {
        return;
    }
    if (!coro_is_main()) {
        // switch current to SUSPEND
        coroSuspend_.push_back(coroReady_.remove(current_->node()));
        current_->setStatus(Status_::SUSPEND);
        current_->waitFor(WaitStatus_::YIELD);
    }
    switchTo_(it->second, Status_::READY);
}

void CoroSched_::cancel(CoroID id) {
    auto it = coros_.find(id);
    if (it == coros_.end()) {
        return;
    }
    it->second->setCancel();
}

void CoroSched_::close(CoroID id) {
    auto it = coros_.find(id);
    if (it != coros_.end()) {
        // Wake up a suspend coroutine and yield cpu to it
        auto* coro = it->second;
        switch (coro->getStatus()) {
            case Status_::READY:
                coroDead_.push_back(coroReady_.remove(coro->node()));
                break;
            case Status_::SUSPEND:
                coroDead_.push_back(coroSuspend_.remove(coro->node()));
                break;
            case Status_::DEAD:
                return;
            default:
                break;
        }
        coro->setStatus(Status_::DEAD);
    }
}

Call& CoroSched_::getRemoteCall() {
    return *remoteCall_;
}

Coro_::Coro_(CoroEntry entry, void* usrPtr) {
    coroID_ = CoroSched_::instance().genCoroID_();
    fiberEntry_ = entry;
    usrPtr_ = usrPtr;
    // Create a windows fiber
    fiber_ = ::CreateFiber(STACK_SIZE, &CoroSched_::entry_, reinterpret_cast<LPVOID>(this));
    // Yield to new fiber, switch status from NONE to READY
    CoroSchedRef.switchTo_(this, Status_::READY);
}

Coro_::~Coro_() {
    if (fiber_ != INVALID_CORO_HANDLE) {
        ::DeleteFiber(fiber_);
    }
    fiber_ = INVALID_CORO_HANDLE;
    status_ = Status_::NONE;
    usrPtr_ = nullptr;
    fiberEntry_ = nullptr;
    coroID_ = INVALID_CORO_ID;
    node_ = nullptr;
    if (timerHandle_) {
        timerHandle_->deleted = true;
        timerHandle_ = nullptr;
    }
}

bool Coro_::isWaitBy(WaitStatus_ sstatus) {
    return waitFlag_[static_cast<std::size_t>(sstatus)];
}

bool Coro_::checkAndReset(WaitStatus_ sstatus) {
    auto flag = isWaitBy(sstatus);
    resetWait(sstatus);
    return flag;
}

void Coro_::waitFor(WaitStatus_ sstatus) {
    waitFlag_[static_cast<std::size_t>(sstatus)] = true;
}

void Coro_::resetWait(WaitStatus_ sstatus) {
    if (sstatus == WaitStatus_::NONE) {
        waitFlag_.reset();
    } else {
        waitFlag_[static_cast<std::size_t>(sstatus)] = false;
    }
}

void Coro_::setCancel() {
    cancel_flag_ = true;
}

bool Coro_::isCancel() {
    return cancel_flag_;
}

}
